package random

import (
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/resolver"
	"log"
	"math/rand/v2"
)

type WeightBuilder struct {
}

func (b *WeightBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	connList := make([]*WeightConn, 0, len(info.ReadySCs))
	totalWeight := 0
	for subConn, sc := range info.ReadySCs {
		weight := sc.Address.Attributes.Value("weight").(int)
		connList = append(connList, &WeightConn{
			conn:      subConn,
			addr:      sc.Address,
			minWeight: totalWeight,
			maxWeight: totalWeight + weight,
		})
		totalWeight += weight
	}

	bl := &WeightBalancer{
		connList:    connList,
		totalWeight: totalWeight,
	}

	return bl
}

type WeightBalancer struct {
	connList    []*WeightConn
	totalWeight int
}

func (b *WeightBalancer) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(b.connList) == 0 {
		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
	}
	// 加权随机

	weight := rand.IntN(b.totalWeight)
	conn := b.connList[0]
	for _, c := range b.connList {
		if weight <= c.maxWeight && weight >= c.minWeight {
			conn = c
		}
	}

	return balancer.PickResult{
		SubConn: conn.conn,
		Done: func(info balancer.DoneInfo) {
			if info.Err != nil {
				log.Println(info.Err)
				return
			}
		},
	}, nil
}

type WeightConn struct {
	conn      balancer.SubConn
	addr      resolver.Address
	minWeight int
	maxWeight int
}
